import os

import numpy as np

from config.config import config
from utils.interval import partition_interval
from utils.read import read_pkl
from utils.write import write_pkl


def main_add_column():
    # 为选择的cluster_count个簇中的每个特征添加attr_interval区间列
    projectPath = config["project_path"]
    dataDirPath = "%s/%s" % (projectPath, config["cluster"]["attrs_select_result_dir"])
    gtDataFileName = os.path.split(config["input"]["gt_data_file_path"])[1].split(".")[0]
    selectMode = config["cluster"]["select_mode"]
    dataPath = r"%s/%s_select_%s.pkl" % (dataDirPath, gtDataFileName, selectMode)

    data = read_pkl(dataPath)
    # 添加列
    print("add column start...")
    interval = partition_interval()  # 区间划分
    intervalList = []
    handler = config["handler"]
    for index, row in data.iterrows():
        attr = row[handler["attr"]]
        i = int(attr / handler["conditions"][2])
        if i < 0 or i >= len(interval):
            intervalList.append(np.NaN)
        else:
            intervalList.append(interval[i])
    data["%s_interval" % handler["attr"]] = [i for i in intervalList]
    savePath = r"%s/%s_select_%s_%s.pkl" % (dataDirPath, gtDataFileName, selectMode, handler["attr"])
    print("添加列后结果存储路径：%s" % savePath)
    write_pkl(savePath, data)


if __name__ == '__main__':
    main_add_column()
